use std::{
    error::Error,
    net::Ipv4Addr,
    net::Ipv6Addr,
};

use clap::Parser;

use libp2p::{
    futures::StreamExt,
    core::{multiaddr::Protocol, Multiaddr}, identify, identity, noise, ping, relay, swarm::{NetworkBehaviour, SwarmEvent}, tcp, yamux
};

use tracing_subscriber::EnvFilter;

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let _ = tracing_subscriber::fmt()
        .with_env_filter(EnvFilter::from_default_env())
        .try_init();
    
    let opt = Opt::parse();
    
    let local_key: identity::Keypair = generate_ed25519(opt.secret_key_seed);
    let peer_id = local_key.public().to_peer_id();
    println!("Relay server PeerID: {}", peer_id);

    let mut swarm = libp2p::SwarmBuilder::with_existing_identity(local_key)
        .with_tokio()
        .with_tcp(
            tcp::Config::default(), 
            noise::Config::new,
            yamux::Config::default,
        )?
        .with_quic()
        .with_behaviour(|key| Behaviour {
            identify: identify::Behaviour::new(identify::Config::new(
                "/TODO/0.0.1".to_string(), key.public())),
            relay: relay::Behaviour::new(key.public().to_peer_id(), Default::default()),
            ping: ping::Behaviour::new(ping::Config::new()),
        })?
        .build();

    let listen_addr_tcp = Multiaddr::empty()
        .with(match opt.use_ipv6 {
            Some(true) => Protocol::from(Ipv6Addr::UNSPECIFIED),
            _ => Protocol::from(Ipv4Addr::UNSPECIFIED),
        })
        .with(Protocol::Tcp(opt.port));

    swarm.listen_on(listen_addr_tcp)?;

    let listen_addr_quic = Multiaddr::empty()
        .with(match opt.use_ipv6 {
            Some(true) => Protocol::from(Ipv6Addr::UNSPECIFIED),
            _ => Protocol::from(Ipv4Addr::UNSPECIFIED),
        })
        .with(Protocol::Udp(opt.port))
        .with(Protocol::QuicV1);
    swarm.listen_on(listen_addr_quic)?;

    // 如果提供了公网IP，构造并广播公网地址
    if let Some(ref public_ip) = opt.public_ip {
        // 构造TCP公网地址
        let public_addr_tcp = Multiaddr::empty()
            .with(Protocol::from(public_ip.parse::<Ipv4Addr>().expect("Invalid public IP address")))
            .with(Protocol::Tcp(opt.port));
        swarm.add_external_address(public_addr_tcp.clone());
        println!("Advertising public TCP address: {public_addr_tcp:?}");
        
        // 构造QUIC公网地址
        let public_addr_quic = Multiaddr::empty()
            .with(Protocol::from(public_ip.parse::<Ipv4Addr>().expect("Invalid public IP address")))
            .with(Protocol::Udp(opt.port))
            .with(Protocol::QuicV1);
        swarm.add_external_address(public_addr_quic.clone());
        println!("Advertising public QUIC address: {public_addr_quic:?}");
    }

    loop {
        match swarm.next().await.expect("Infinite Stream") {
            SwarmEvent::Behaviour(event) => {
                if let BehaviourEvent::Identify(identify::Event::Received { 
                    info: identify::Info { observed_addr, .. },
                    ..
                }) = &event
                {
                    swarm.add_external_address(observed_addr.clone());
                }
            
                println!("{event:?}")
            }
            SwarmEvent::NewListenAddr { address, .. } => {
                println!("Listening on {address:?}");
            }
            _ => {}
        }
    }
}

#[derive(NetworkBehaviour)]
struct Behaviour {
    identify: identify::Behaviour,
    relay: relay::Behaviour,
    ping: ping::Behaviour,
}

fn generate_ed25519(secret_key_seed: u8) -> identity::Keypair {
    let mut bytes = [0u8; 32];
    bytes[0] = secret_key_seed;
    identity::Keypair::ed25519_from_bytes(bytes).expect("only errors on wrong length")
}

#[derive(Parser, Debug)]
#[command(name = "my_relay_server")]
struct Opt {
    #[arg(long)]
    use_ipv6: Option<bool>,

    #[arg(long)]
    secret_key_seed: u8,

    #[arg(long)]
    port: u16,
    
    #[arg(long, help = "Public IP address to advertise to clients")]
    public_ip: Option<String>,
}
